{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "720746b5",
   "metadata": {},
   "source": [
    "# 自然语言推断：微调BERT\n",
    ":label:`sec_natural-language-inference-bert`\n",
    "\n",
    "在本章的前面几节中，我们已经为SNLI数据集（ :numref:`sec_natural-language-inference-and-dataset`）上的自然语言推断任务设计了一个基于注意力的结构（ :numref:`sec_natural-language-inference-attention`）。现在，我们通过微调BERT来重新审视这项任务。正如在 :numref:`sec_finetuning-bert`中讨论的那样，自然语言推断是一个序列级别的文本对分类问题，而微调BERT只需要一个额外的基于多层感知机的架构，如 :numref:`fig_nlp-map-nli-bert`中所示。\n",
    "\n",
    "![将预训练BERT提供给基于多层感知机的自然语言推断架构](../img/nlp-map-nli-bert.svg)\n",
    ":label:`fig_nlp-map-nli-bert`\n",
    "\n",
    "本节将下载一个预训练好的小版本的BERT，然后对其进行微调，以便在SNLI数据集上进行自然语言推断。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02cd4c5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5ec3d66",
   "metadata": {},
   "outputs": [],
   "source": [
    "from d2l import mindspore as d2l\n",
    "import json\n",
    "import multiprocessing\n",
    "from mindspore import nn, ops\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e150d09",
   "metadata": {},
   "source": [
    "## [**加载预训练的BERT**]\n",
    "\n",
    "我们已经在 :numref:`sec_bert-dataset`和 :numref:`sec_bert-pretraining`WikiText-2数据集上预训练BERT（请注意，原始的BERT模型是在更大的语料库上预训练的）。正如在 :numref:`sec_bert-pretraining`中所讨论的，原始的BERT模型有数以亿计的参数。在下面，我们提供了两个版本的预训练的BERT：“bert.base”与原始的BERT基础模型一样大，需要大量的计算资源才能进行微调，而“bert.small”是一个小版本，以便于演示。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d061884d",
   "metadata": {},
   "outputs": [],
   "source": [
    "d2l.DATA_HUB['bert_small'] = ('https://authoring-modelarts-cnsouth222.cloudbrain2.pcl.ac.cn/32d9fa25-4f49-4173-a958-0533932fb920/files/bert.small.ms.zip?_xsrf=2%7C52321e82%7C6ce0f84768db47601fa38e4b5c803120%7C1677304596',\n",
    "'0daa723ea0d1610c20ddc7d34427aa9e')\n",
    "\n",
    "d2l.DATA_HUB['bert_base'] = ('https://authoring-modelarts-cnsouth222.cloudbrain2.pcl.ac.cn/32d9fa25-4f49-4173-a958-0533932fb920/files/bert.base.ms.zip?_xsrf=2%7C52321e82%7C6ce0f84768db47601fa38e4b5c803120%7C1677304596',\n",
    "'96f5049ad083860728acc1f475a47bf0')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ec3a368",
   "metadata": {},
   "source": [
    "两个预训练好的BERT模型都包含一个定义词表的“vocab.json”文件和一个预训练参数的“pretrained.params”文件。我们实现了以下`load_pretrained_model`函数来[**加载预先训练好的BERT参数**]。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76863b70",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_pretrained_model(pretrained_model, num_hiddens, ffn_num_hiddens,\n",
    "                          num_heads, num_layers, dropout, max_len):\n",
    "    data_dir = \"../data\"\n",
    "    # 定义空词表以加载预定义词表\n",
    "    vocab = d2l.Vocab()\n",
    "    vocab.idx_to_token = json.load(open(os.path.join(data_dir,\n",
    "        'vocab.json')))\n",
    "    vocab.token_to_idx = {token: idx for idx, token in enumerate(\n",
    "        vocab.idx_to_token)}\n",
    "    bert = d2l.BERTModel(len(vocab), num_hiddens, norm_shape=[256],\n",
    "                         ffn_num_input=256, ffn_num_hiddens=ffn_num_hiddens,\n",
    "                         num_heads=4, num_layers=2, dropout=0.2,\n",
    "                         max_len=max_len, key_size=256, query_size=256,\n",
    "                         value_size=256, hid_in_features=256,\n",
    "                         mlm_in_features=256, nsp_in_features=256)\n",
    "\n",
    "    # 加载预训练BERT参数\n",
    "    param_dict = load_checkpoint(os.path.join(data_dir,\n",
    "                                             'pretrained.ckpt'))\n",
    "    load_param_into_net(bert, param_dict)\n",
    "    return bert, vocab"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ab0cd17",
   "metadata": {},
   "source": [
    "为了便于在大多数机器上演示，我们将在本节中加载和微调经过预训练BERT的小版本（“bert.small”）。在练习中，我们将展示如何微调大得多的“bert.base”以显著提高测试精度。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "714628e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "bert, vocab = load_pretrained_model(\n",
    "    'bert_small', num_hiddens=256, ffn_num_hiddens=512, num_heads=4,\n",
    "    num_layers=2, dropout=0.1, max_len=512)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "716ff657",
   "metadata": {},
   "source": [
    "## [**微调BERT的数据集**]\n",
    "\n",
    "对于SNLI数据集的下游任务自然语言推断，我们定义了一个定制的数据集类`SNLIBERTDataset`。在每个样本中，前提和假设形成一对文本序列，并被打包成一个BERT输入序列，如 :numref:`fig_bert-two-seqs`所示。回想 :numref:`subsec_bert_input_rep`，片段索引用于区分BERT输入序列中的前提和假设。利用预定义的BERT输入序列的最大长度（`max_len`），持续移除输入文本对中较长文本的最后一个标记，直到满足`max_len`。为了加速生成用于微调BERT的SNLI数据集，我们使用4个工作进程并行生成训练或测试样本。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3cbc740",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SNLIBERTDataset:\n",
    "    def __init__(self, dataset, max_len, vocab=None):\n",
    "        all_premise_hypothesis_tokens = [[\n",
    "            p_tokens, h_tokens] for p_tokens, h_tokens in zip(\n",
    "            *[d2l.tokenize([s.lower() for s in sentences])\n",
    "              for sentences in dataset[:2]])]\n",
    "        \n",
    "        self.labels = np.array(dataset[2])\n",
    "        self.vocab = vocab\n",
    "        self.max_len = max_len\n",
    "        (self.all_token_ids, self.all_segments,\n",
    "         self.valid_lens) = self._preprocess(all_premise_hypothesis_tokens)\n",
    "        print('read ' + str(len(self.all_token_ids)) + ' examples')\n",
    "\n",
    "    def _preprocess(self, all_premise_hypothesis_tokens):\n",
    "        pool = multiprocessing.Pool(4)  # 使用4个进程\n",
    "        out = pool.map(self._mp_worker, all_premise_hypothesis_tokens)\n",
    "        all_token_ids = [\n",
    "            token_ids for token_ids, segments, valid_len in out]\n",
    "        all_segments = [segments for token_ids, segments, valid_len in out]\n",
    "        valid_lens = [valid_len for token_ids, segments, valid_len in out]\n",
    "        return (np.array(all_token_ids, dtype=np.long),\n",
    "                np.array(all_segments, dtype=np.long), \n",
    "                np.array(valid_lens))\n",
    "\n",
    "    def _mp_worker(self, premise_hypothesis_tokens):\n",
    "        p_tokens, h_tokens = premise_hypothesis_tokens\n",
    "        self._truncate_pair_of_tokens(p_tokens, h_tokens)\n",
    "        tokens, segments = d2l.get_tokens_and_segments(p_tokens, h_tokens)\n",
    "        token_ids = self.vocab[tokens] + [self.vocab['<pad>']] \\\n",
    "                             * (self.max_len - len(tokens))\n",
    "        segments = segments + [0] * (self.max_len - len(segments))\n",
    "        valid_len = len(tokens)\n",
    "        return token_ids, segments, valid_len\n",
    "\n",
    "    def _truncate_pair_of_tokens(self, p_tokens, h_tokens):\n",
    "        # 为BERT输入中的'<CLS>'、'<SEP>'和'<SEP>'词元保留位置\n",
    "        while len(p_tokens) + len(h_tokens) > self.max_len - 3:\n",
    "            if len(p_tokens) > len(h_tokens):\n",
    "                p_tokens.pop()\n",
    "            else:\n",
    "                h_tokens.pop()\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return (self.all_token_ids[idx], self.all_segments[idx],\n",
    "                self.valid_lens[idx]), self.labels[idx]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.all_token_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ede7f40",
   "metadata": {},
   "source": [
    "下载完SNLI数据集后，我们通过实例化`SNLIBERTDataset`类来[**生成训练和测试样本**]。这些样本将在自然语言推断的训练和测试期间进行小批量读取。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d000c29",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size, max_len, num_workers = 512, 128, d2l.get_dataloader_workers()\n",
    "data_dir = d2l.download_extract('SNLI')\n",
    "train_set = SNLIBERTDataset(d2l.read_snli(data_dir, True), max_len, vocab)\n",
    "test_set = SNLIBERTDataset(d2l.read_snli(data_dir, False), max_len, vocab)\n",
    "train_iter = ds.GeneratorDataset(train_set, shuffle=True, num_parallel_workers=num_workers, column_names=[\"0\", \"1\", \"2\", \"3\"])\n",
    "train_iter = train_iter.batch(batch_size=batch_size)\n",
    "test_iter = ds.GeneratorDataset(test_set, num_parallel_workers=num_workers, column_names=[\"0\", \"1\", \"2\", \"3\"])\n",
    "test_iter = test_iter.batch(batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c0a3284",
   "metadata": {},
   "source": [
    "## 微调BERT\n",
    "\n",
    "如 :numref:`fig_bert-two-seqs`所示，用于自然语言推断的微调BERT只需要一个额外的多层感知机，该多层感知机由两个全连接层组成（请参见下面`BERTClassifier`类中的`self.hidden`和`self.output`）。[**这个多层感知机将特殊的“&lt;cls&gt;”词元**]的BERT表示进行了转换，该词元同时编码前提和假设的信息(**为自然语言推断的三个输出**)：蕴涵、矛盾和中性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "186e8633",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BERTClassifier(nn.Cell):\n",
    "    def __init__(self, bert):\n",
    "        super(BERTClassifier, self).__init__()\n",
    "        self.encoder = bert.encoder\n",
    "        self.hidden = bert.hidden\n",
    "        self.output = nn.Dense(256, 3)\n",
    "\n",
    "    def construct(self, inputs):\n",
    "        tokens_X, segments_X, valid_lens_x = inputs\n",
    "        encoded_X = self.encoder(tokens_X, segments_X, valid_lens_x.squeeze(1))\n",
    "        return self.output(self.hidden(encoded_X[:, 0, :]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a89385cb",
   "metadata": {},
   "source": [
    "在下文中，预训练的BERT模型`bert`被送到用于下游应用的`BERTClassifier`实例`net`中。在BERT微调的常见实现中，只有额外的多层感知机（`net.output`）的输出层的参数将从零开始学习。预训练BERT编码器（`net.encoder`）和额外的多层感知机的隐藏层（`net.hidden`）的所有参数都将进行微调。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ef54314",
   "metadata": {},
   "outputs": [],
   "source": [
    "net = BERTClassifier(bert)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5201ff6",
   "metadata": {},
   "source": [
    "回想一下，在 :numref:`sec_bert`中，`MaskLM`类和`NextSentencePred`类在其使用的多层感知机中都有一些参数。这些参数是预训练BERT模型`bert`中参数的一部分，因此是`net`中的参数的一部分。然而，这些参数仅用于计算预训练过程中的遮蔽语言模型损失和下一句预测损失。这两个损失函数与微调下游应用无关，因此当BERT微调时，`MaskLM`和`NextSentencePred`中采用的多层感知机的参数不会更新（陈旧的，staled）。\n",
    "\n",
    "为了允许具有陈旧梯度的参数，标志`ignore_stale_grad=True`在`step`函数`d2l.train_batch_ch13`中被设置。我们通过该函数使用SNLI的训练集（`train_iter`）和测试集（`test_iter`）对`net`模型进行训练和评估。由于计算资源有限，[**训练**]和测试精度可以进一步提高：我们把对它的讨论留在练习中。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc884fb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr, num_epochs = 1e-4, 5\n",
    "trainer = nn.Adam(learning_rate=lr, params=net.trainable_params())\n",
    "loss = nn.CrossEntropyLoss(reduction=\"none\")\n",
    "d2l.train_ch15_bert(net, train_iter, test_iter, loss, trainer, num_epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56596a5c",
   "metadata": {},
   "source": [
    "## 小结\n",
    "\n",
    "* 我们可以针对下游应用对预训练的BERT模型进行微调，例如在SNLI数据集上进行自然语言推断。\n",
    "* 在微调过程中，BERT模型成为下游应用模型的一部分。仅与训练前损失相关的参数在微调期间不会更新。\n",
    "\n",
    "## 练习\n",
    "\n",
    "1. 如果您的计算资源允许，请微调一个更大的预训练BERT模型，该模型与原始的BERT基础模型一样大。修改`load_pretrained_model`函数中的参数设置：将“bert.small”替换为“bert.base”，将`num_hiddens=256`、`ffn_num_hiddens=512`、`num_heads=4`和`num_layers=2`的值分别增加到768、3072、12和12。通过增加微调迭代轮数（可能还会调优其他超参数），你可以获得高于0.86的测试精度吗？\n",
    "1. 如何根据一对序列的长度比值截断它们？将此对截断方法与`SNLIBERTDataset`类中使用的方法进行比较。它们的利弊是什么？"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
